import torch
import torch.nn as nn
# 定义输入数据：序列长度seq_len=5,批次大小batch_size=3,特征维度input_size=10
input = torch.randn(5, 3, 10) # (seq_len,batch_Size,input_size)

# 输入input序列长度seq_len=5, 每次取minibatch大小,
# batch_size=3, 数据向量维数=10（为x的维度）。
# 每次运行时取3个含有5个字的句子（且句子中每个字的维度为10进行运行）
# 输入数据x的向量维数10, 设定lstm隐藏层的特征维度20,2个lstm层。如果是1,可以省略
lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)

# 初始化的隐藏元和记忆元,通常它们的维度是一样的
# 2个LSTM层,batch_size=3, 隐藏层的特征维度20 ,output是最后一层lstm的每个词向量对应隐藏层输出
h_0 = torch.randn(2, 3, 20) # (num_layers,batch_size,hidden_size) 
c_0 = torch.randn(2, 3, 20) # (num_layers, batch_size, hidden_size)

# 前向传播
# h_n,c_n是所有层最后一个隐藏元和记忆元的输出
output, (h_n, c_n) = lstm(input, (h_0, c_0)) 
# 如果未提供（h_0,c_0）,则默认为零。
# 由于batch_first=False  
# output:(seq_len,batch_size,hidden_size) 
# h_n,c_n: (num_layers,batch_size,hidden_size) (2, 3, 20) 

# 输出形状说明
print(output.size(), h_n.size(), c_n.size())  
